"""Solvers for the 1d Euler equations"""

# base class for all Solvers
import solver

from optparse import OptionGroup
from os.path import join
import numpy

from hogs.solvers.flux.euler.lagrangian_flux import flux_types

#############################################################################
# `EESolver` class
#############################################################################
class EESolver(solver.Solver1D):
    """Base class for all solvers for the Euler Equations (EE)""" 

    def __init__(self, gamma, nvar, tf, grid=None):
        """Constructor for the solver.

        Parameters:
        -----------

        gamma : float
            Ratio of specific heats. Constant

        nvar : int
            The number of conserved variables.

        tf : float
            Final time

        cfl : float
            CFL constant

        pfreq : int
            Output print frequency.

        """
        solver.Solver1D.__init__(self, nvar, tf, grid=grid)

        self.gamma = gamma
        self.flux_function = None

    def _check_flux_function(self):
        if self.flux_function is None:
            msg = """This solver's flux function is None. You must specify a
            valid flux function from fvmhslc.solvers.flux_functions
            """
            raise RuntimeError(msg)

###############################################################################
# `EELagrangianSolver` class
###############################################################################
class EELagrangianSolver(EESolver):
    """Lagrangian type scheme of Shu"""

    def __init__(self, gamma, nvar, tf, flux_type=0, grid=None):
        EESolver.__init__(self,gamma=gamma, nvar=nvar, tf=tf, grid=grid)
        self.defaults = dict(flux_type=flux_type)

    def compute_dt(self):
        """Compute the time step.

        The time step criterion uses is:

        ..math::

        \Delta t \leq \lambda min_i( \frac{\Delta x_i}{a_i} )

        The Courant number $\lambda$ is equal to 0.6
        
        """

        grid = self.grid

        # get the mesh points and the sound speed
        x = grid.x
        a = grid.a

        # get the cell sizes
        dx =  x[1:] - x[:-1]

        return 2.5e-7
        #return self.cfl * min( dx/a )

    def update(self, dt):
        """Advance the solution"""

        grid = self.grid

        # get the face coordinates
        x = grid.x
        
        # compute the flux
        self.flux_function.compute_flux(grid, dt)

        # now update the cell values
        for cell in grid.cells[grid.nb:-grid.nb]:

            # get the flux at the left and right faces
            fl, fr = cell.get_fluxes()

            # get the cell lengths in the previous time step
            dxold = cell.dx
            dxnew = dxold + (cell.right.vertex_vel - cell.left.vertex_vel) * dt
            cell.dx = dxnew

            # update the cell variables
            q = cell.q
            q[:] = (q * dxold - dt * (fr - fl))/dxnew

         # now update the mesh points
        for face in grid.faces:
            face.update_position(dt)

    def get_options(self, parser):
        """Set options for the Lagrangian solver"""

        # option for the flux function
        opt = OptionGroup(parser, "Options for the Lagrangian flux")
        opt.add_option("--flux", action="store", type="int",
                       dest="flux_type", default=self.defaults["flux_type"],
                       help="The lagrangian flux function to use\t\t\t" + \
                       "".join(['%d - %-51s\t\t'%(d,s[0]) for d,s in
                                 enumerate(flux_types)]))

        # optional options for specific flux functions. We might want
        # to define a general set of options for all lagrangian flux
        # functions which may be implemented for the specific types.
        #TODO

        return opt
                       
    def set_options(self, options):
        """Set the options"""
        options = options or self.defaults

        # set the flux function for the solver
        flux_type = options.get("flux_type")
        self.flux_function = flux_types[flux_type][1]()
